#ifdef WIN32
#include <windows.h>
#include <tlhelp32.h>
#include <Psapi.h>
#endif
#include <sys/stat.h>
#include <process.h>
#include <fcntl.h>
#include <stdio.h>
#include <fcntl.h>
#include <io.h>
#include <stdlib.h>
#include <string.h>
#include <sys/types.h>
#include "processwrapper.h"

#define SF_BUFFER_SIZE 256

CProcessWrapper::CProcessWrapper(int flags):envsize(0),pchCommandLine(0),conbuf(0),hProcess(0)
{
	this->flags = flags;
#ifdef _USE_WIN32_PROC
	hStdinWrite = 0;
	hStdoutRead = 0;
	hStderrRead = 0;
#else
	fdStdinWrite = 0;
	fdStdoutRead = 0;
	fdStderrRead = 0;
#endif
}

int CProcessWrapper::ReadStdout()
{
#ifdef _USE_WIN32_PROC
	return ReadOutput(hStdoutRead);
#else
	return ReadOutput(fdStdoutRead);
#endif
}

int CProcessWrapper::ReadStderr()
{
#ifdef _USE_WIN32_PROC
	return ReadOutput(hStderrRead);
#else
	return ReadOutput(fdStderrRead);
#endif
}

void CProcessWrapper::Cleanup()
{
	CloseStdin();
	CloseStdout();
	CloseStderr();
	if (IsProcessRunning()) {
		Terminate();
	}
	FreeBuffer();
	//CloseHandle(hProcess);
}

int CProcessWrapper::WaitIdle(int milliseconds)
{
	return WaitForInputIdle(hProcess, milliseconds);
}

typedef DWORD (WINAPI *PFNGetProcessId)(HANDLE hProcess);
typedef BOOL (WINAPI *PFNEnumProcesses)(DWORD* pProcessIds, DWORD cb, DWORD* pBytesReturned);
typedef HANDLE (WINAPI *PFNOpenThread)(DWORD dwDesiredAccess, BOOL bInheritHandle, DWORD dwProcessId);

static HMODULE hKernel = LoadLibrary("kernel32.dll");
static HMODULE hPsapi = LoadLibrary("Psapi.dll");
static PFNGetProcessId pfnGetPid = (PFNGetProcessId)GetProcAddress(hKernel, "GetProcessId");
static PFNOpenThread pfnOpenThread = (PFNOpenThread)GetProcAddress(hKernel, "OpenThread");
static PFNEnumProcesses pfnEnumProc = (PFNEnumProcesses)GetProcAddress(hPsapi, "EnumProcesses");

int CProcessWrapper::GetProcessIdByHandle(HANDLE hProcess)
{
	if (pfnGetPid) {
		return (*pfnGetPid)(hProcess);
	} else if (pfnEnumProc) {
		DWORD dwProcessID = -1;
		DWORD aProcesses[256], cbNeeded, cProcesses;
		unsigned int i;

		if (!(*pfnEnumProc)( aProcesses, sizeof(aProcesses), &cbNeeded ) ) return -1;

		// Calculate how many process identifiers were returned.
		cProcesses = cbNeeded / sizeof(DWORD);

		// Print the name and process identifier for each process.
		for ( i = 0; i < cProcesses; i++ ) {
			HANDLE hTempProcess = OpenProcess( PROCESS_QUERY_INFORMATION, FALSE, aProcesses[i] );
			if (!hTempProcess) continue;
			CloseHandle( hTempProcess );
			if(hTempProcess == hProcess) {
				dwProcessID = aProcesses[i];
			}
		}
		return dwProcessID;
	}
	return 0;
}

bool CProcessWrapper::Spawn(char* commandLine, char* curDir, char** env)
{
#ifdef _USE_WIN32_PROC
	BOOL fSuccess; 
	PROCESS_INFORMATION piProcInfo;
	SECURITY_ATTRIBUTES saAttr;
	HANDLE hChildStdoutRd = 0, hChildStdoutWr = 0;
	HANDLE hChildStderrRd = 0, hChildStderrWr = 0;
	HANDLE hChildStdinRd = 0, hChildStdinWr = 0;
	saAttr.nLength = sizeof(SECURITY_ATTRIBUTES); 
	saAttr.bInheritHandle = TRUE; 
	saAttr.lpSecurityDescriptor = NULL; 
#else
	int fdStdinRead = 0, fdStdinOld = 0;
	int fdStdoutWrite = 0, fdStdoutOld = 0;
	int fdStderrWrite = 0, fdStderrOld = 0;
#endif

	Cleanup();

	STARTUPINFO siStartInfo;
	memset(&siStartInfo, 0, sizeof(siStartInfo));

	if (flags & SF_REDIRECT_STDIN) {
#ifndef _USE_WIN32_PROC
		MakePipe(fdStdinRead, fdStdinWrite, 16384, false, false);
		//fdStdinOld = _dup(0);
		_dup2(fdStdinRead, 0);
#else
	CreatePipe(&hChildStdinRd, &hChildStdinWr, &saAttr, 0);
	SetHandleInformation( hChildStdinWr, HANDLE_FLAG_INHERIT, 0);
	siStartInfo.hStdInput = hChildStdinRd;
	siStartInfo.dwFlags |= STARTF_USESTDHANDLES;
#endif
	}
	if (flags & SF_REDIRECT_STDOUT) {
#ifndef _USE_WIN32_PROC
		MakePipe(fdStdoutRead, fdStdoutWrite, 16384, false, false);
		fdStdoutOld = _dup(1);
		_dup2(fdStdoutWrite, 1);
#else
		CreatePipe(&hChildStdoutRd, &hChildStdoutWr, &saAttr, 0);
		SetHandleInformation(hChildStdoutRd, HANDLE_FLAG_INHERIT, 0);
		siStartInfo.hStdOutput = hChildStdoutWr;
		siStartInfo.dwFlags |= STARTF_USESTDHANDLES;
#endif
		if ((flags & SF_REDIRECT_OUTPUT) == SF_REDIRECT_OUTPUT) {
#ifndef _USE_WIN32_PROC
			fdStderrOld = _dup(2);
			_dup2(fdStdoutWrite, 2);
#else
			siStartInfo.hStdError = siStartInfo.hStdOutput;
#endif
		}
	}
	if (flags & SF_REDIRECT_STDERR) {
#ifndef _USE_WIN32_PROC
		MakePipe(fdStderrRead, fdStderrWrite, 16384, false, false);
		fdStderrOld = _dup(2);
		_dup2(fdStderrWrite, 2);
		
#else
		CreatePipe(&hChildStderrRd, &hChildStderrWr, &saAttr, 0);
		SetHandleInformation(hChildStderrWr, HANDLE_FLAG_INHERIT, 0);
		siStartInfo.hStdError = hChildStderrWr;
		siStartInfo.dwFlags |= STARTF_USESTDHANDLES;
#endif
	}

#ifndef _USE_WIN32_PROC
	char *pchCommandLine = _strdup(commandLine);
	char *oldCurDir = 0;
	char **tokens;
	char *exe;
	tokens = Tokenize(pchCommandLine);
	if (tokens[0][0] == '\"') {
		exe = _strdup(tokens[0] + 1);
		char *p = strrchr(exe, '\"');
		if (p) *p = 0;
	} else {
		exe = _strdup(tokens[0]);
	}
	if (curDir) {
		oldCurDir = (char*)malloc(MAX_PATH);
		GetCurrentDirectory(MAX_PATH, oldCurDir);
		SetCurrentDirectory(curDir);
	}
	if (env) {
		hProcess = (HANDLE)_spawnvpe( _P_NOWAIT, exe, tokens, env);
	} else {
		hProcess = (HANDLE)_spawnvp( _P_NOWAIT, exe, tokens);
	}
	pid = GetProcessIdByHandle(hProcess);

	if (oldCurDir) {
		SetCurrentDirectory(oldCurDir);
		free(oldCurDir);
	}
	free(exe);
	free(tokens);
	free(pchCommandLine);
#else

	// Set up members of the STARTUPINFO structure.
	siStartInfo.cb = sizeof(STARTUPINFO);
	siStartInfo.wShowWindow = SW_HIDE;
	siStartInfo.dwFlags |= STARTF_USESHOWWINDOW;
	if (flags & (SF_REDIRECT_STDIN | SF_REDIRECT_STDOUT | SF_REDIRECT_STDERR)) siStartInfo.dwFlags |= STARTF_USESTDHANDLES;

	fSuccess = CreateProcess(0,
			commandLine,				// command line
			NULL,					// process security attributes
			NULL,					// primary thread security attributes
			TRUE,					// handles are inherited
			0,						// creation flags
			NULL,					// process' environment
			curDir,					// current directory
			&siStartInfo,			// STARTUPINFO pointer
			&piProcInfo);			// receives PROCESS_INFORMATION

	if (fSuccess) {
		hProcess = piProcInfo.hProcess;
		pid = piProcInfo.dwProcessId;
	}
	//WaitIdle(3000);
#endif

#ifdef _USE_WIN32_PROC
	if (hChildStdoutWr) CloseHandle(hChildStdoutWr);
	if (hChildStdinRd) CloseHandle(hChildStdinRd);
	hStdinWrite = hChildStdinWr;
	hStdoutRead = hChildStdoutRd;
	hStderrRead = hChildStderrRd;
#else
	if (fdStdinRead) _close(fdStdinRead);
	if (fdStdoutWrite) _close(fdStdoutWrite);
	if (fdStderrWrite) _close(fdStderrWrite);
	//if (fdStdinOld) _dup2( fdStdinOld, 0);
	if (fdStdoutOld) _dup2( fdStdoutOld, 1 );
	if (fdStderrOld) _dup2( fdStderrOld, 2 );
#endif
	if (!hProcess || hProcess == INVALID_HANDLE_VALUE) {
		hProcess = 0;
		Cleanup();
		return false;
	}

	return true;
}

int CProcessWrapper::Run(char* commandLine, char* curDir)
{
	char *oldCurDir = 0;
	char *cmdline = _strdup(commandLine);
	char **tokens;
	char *exe;
	tokens = Tokenize(cmdline);
	if (commandLine[0] == '\"') {
		char *end = strchr(commandLine + 1, '\"');
		int len = end - commandLine - 1;
		exe = (char*)malloc(len + 1);
		memcpy(exe, commandLine + 1, len);
		exe[len] = 0;
	} else {
		exe = _strdup(tokens[0]);
	}
	if (curDir) {
		oldCurDir = (char*)malloc(MAX_PATH);
		GetCurrentDirectory(MAX_PATH, oldCurDir);
		SetCurrentDirectory(curDir);
	}
	int ret = _spawnvp( _P_WAIT, exe, tokens );
	if (oldCurDir) {
		SetCurrentDirectory(oldCurDir);
	}
	free(exe);
	free(tokens);
	free(cmdline);
	return ret;
}

int CProcessWrapper::ReadOutput
#ifdef _USE_WIN32_PROC
	(HANDLE handle)
#else
	(int fd)
#endif
{
	if (!conbuf || conbufsize <= 0) {
		if (flags & SF_ALLOC) {
			conbuf = (char*)malloc(SF_BUFFER_SIZE);
			conbufsize = SF_BUFFER_SIZE;
		} else {
			return -1;
		}
	}

	int offset=0;
	for(;;) {
		int bytes;
		if (offset >= conbufsize - 1) {
			if (flags & SF_ALLOC) {
				conbufsize <<= 1;
				conbuf = (char*)realloc(conbuf, conbufsize);
			} else {
				break;
			}
		}
#ifdef _USE_WIN32_PROC
		if (!ReadFile(handle, conbuf + offset, conbufsize - 1 - offset, (LPDWORD)&bytes, 0)) break;
#else
		bytes = _read(fdStdoutRead, conbuf + offset, conbufsize - 1 - offset);
		if (bytes <= 0) break;
#endif
		offset += bytes;
	}
	conbuf[offset]=0;
	return offset;
}

int CProcessWrapper::Read(void* buffer, int bufsize)
{
#ifdef _USE_WIN32_PROC
	DWORD bytes;
	return (hStdoutRead > 0 && ReadFile(hStdoutRead, buffer, bufsize, &bytes, 0)) ? bytes : -1;
#else
	if (fdStdoutRead <= 0) return -1;
	return _read((int)fdStdoutRead, buffer, bufsize);
#endif
}

int CProcessWrapper::Write(void* data, int len)
{
#ifdef _USE_WIN32_PROC
	DWORD bytes;
	return (hStdinWrite > 0 && WriteFile(hStdinWrite, data, len, &bytes, 0)) ? bytes : -1;
#else
	if (fdStdinWrite <= 0) return -1;
	return _write(fdStdinWrite, data, len);
#endif
}

void CProcessWrapper::AllocateBuffer(int bytes)
{
	if (conbuf) free(conbuf);
	conbuf = (char*)malloc(bytes);
	conbufsize = bytes;
}

void CProcessWrapper::FreeBuffer()
{
	if (conbuf) {
		free(conbuf);
		conbuf = 0;
	}
}

int CProcessWrapper::Wait(int timeout)
{
	switch (WaitForSingleObject(hProcess, timeout)) {
	case WAIT_TIMEOUT:
		return 0;
	case WAIT_OBJECT_0:
		return 1;
	default:
		return -1;
	}
}

bool CProcessWrapper::Terminate()
{
	return TerminateProcess(hProcess,0) ? true : false;
}

void CProcessWrapper::AddEnvVariable(char* envName, char* envValue)
{
	/*
	int varbytes = strlen(envName) + strlen(envValue) + 2;
	pchEnv = (char*)realloc(pchEnv, envsize + varbytes + 1);
	sprintf(pchEnv + envsize, "%s=%s", envName, envValue);
	envsize += varbytes;
	pchEnv[envsize] = 0;
	*/
}

void CProcessWrapper::PurgeVariable()
{
	/*
	if (pchEnv) {
		free(pchEnv);
	}
	*/
}

bool CProcessWrapper::IsProcessRunning(int* pexitcode)
{
	DWORD exitcode = 0;
	if (!hProcess || !GetExitCodeProcess(hProcess, &exitcode)) {
		if (pexitcode) *pexitcode = 0;
		return false;
	}
	if (exitcode == STILL_ACTIVE) return true;
	if (pexitcode) *pexitcode = exitcode;
	return false;

}

bool CProcessWrapper::Suspend(bool resume)
{
	if (!pfnOpenThread || !pid) return false;

    HANDLE        hThreadSnap = NULL; 
    bool          bRet        = false; 
    THREADENTRY32 te32        = {0}; 
 
    // Take a snapshot of all threads currently in the system. 

    hThreadSnap = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0); 
    if (hThreadSnap == INVALID_HANDLE_VALUE) 
        return (FALSE); 
 
    // Fill in the size of the structure before using it. 

    te32.dwSize = sizeof(THREADENTRY32); 
 
    // Walk the thread snapshot to find all threads of the process. 
    // If the thread belongs to the process, add its information 
    // to the display list.
 
    if (Thread32First(hThreadSnap, &te32)) { 
        do 
        { 
            if (te32.th32OwnerProcessID == pid) 
            {
				HANDLE hThread = (*pfnOpenThread)(THREAD_SUSPEND_RESUME, FALSE, te32.th32ThreadID);
				if (hThread) {
					if (resume)
						ResumeThread(hThread);
					else
						SuspendThread(hThread);
					CloseHandle(hThread);
				}
            } 
        }
        while (Thread32Next(hThreadSnap, &te32)); 
        bRet = true; 
    } 
 
    // Do not forget to clean up the snapshot object. 
    CloseHandle (hThreadSnap); 
 
    return (bRet); 
}

bool CProcessWrapper::SetPriority(int priority)
{
	return SetPriorityClass(hProcess, priority) ? true : false;
}

bool CProcessWrapper::SetAffinity(int mask)
{
	return SetProcessAffinityMask(hProcess, mask) ? true : false;
}

#define READ_FD 0
#define WRITE_FD 1

bool CProcessWrapper::MakePipe(int& readpipe, int& writepipe, int bufsize, bool inheritRead, bool inheritWrite, int flags)
{
	int fdpipe[2];
	if (_pipe(fdpipe, bufsize, flags == -1 ? O_NOINHERIT|O_BINARY : flags)) return false;
	if (inheritRead) {
		readpipe = _dup(fdpipe[READ_FD]);
		_close(fdpipe[READ_FD]);
	} else {
		readpipe = fdpipe[READ_FD];
	}
	if (inheritWrite) {
		writepipe = _dup(fdpipe[WRITE_FD]);
		_close(fdpipe[WRITE_FD]);
	} else {
		writepipe = fdpipe[WRITE_FD];
	}
	return true;
}

char** CProcessWrapper::Tokenize(char* str, char delimiter)
{
	char** tokens;
	int n = 1;
	int i;
	char *p;
	
	// find out number of tokens
	p = str;
	for (;;) {
		while (*p && *p != delimiter) p++;
		if (!*p) break;
		n++;
		while (*(++p) == delimiter);
	}
	// allocate buffer for array
	tokens = (char**)malloc((n + 1) * sizeof(char*));
	// store pointers to tokens
	p = str;
	for (i = 0; i < n; i++) {
		tokens[i] = p;
		while (*p && *p != delimiter) p++;
		if (!*p) break;
		*p = 0;
		while (*(++p) == delimiter);
	}
	tokens[n] = 0;
	return tokens;
}

void CProcessWrapper::CloseStdout()
{
#ifdef _USE_WIN32_PROC
	if (hStdoutRead) {
		CloseHandle(hStdoutRead);
		hStdoutRead = 0;
	}
#else
	int fd = fdStdoutRead;
	fdStdoutRead = 0;
	if (fd > 0 ) _close(fd);
#endif
}

void CProcessWrapper::CloseStdin()
{
#ifdef _USE_WIN32_PROC
	if (hStdinWrite) {
		CloseHandle(hStdinWrite);
		hStdinWrite = 0;
	}
#else
	int fd = fdStdinWrite;
	fdStdinWrite = 0;
	if (fd > 0 ) _close(fd);
#endif
}

void CProcessWrapper::CloseStderr()
{
#ifdef _USE_WIN32_PROC
	if (hStderrRead) {
		CloseHandle(hStderrRead);
		hStderrRead = 0;
	}
#else
	int fd = fdStderrRead;
	fdStderrRead = 0;
	if (fd > 0 ) _close(fd);
#endif
}

